#!/usr/bin/python

import roslib; roslib.load_manifest('framer')
import rospy,smach,tf
from move_base_msgs.msg import MoveBaseGoal
from geometry_msgs.msg import Point, PoseStamped,Quaternion, PointStamped
from smach_ros import ActionServerWrapper
from math import sqrt, tan, pi,atan

class TooFarToFrame(Exception):
    pass
class SubjectsTooClose(Exception):
    pass

class RuleOfThirds(smach.State):
    angleOfView = pi/2.0
    minDistance = 1.5
    maxDistance = 3.0
    sensorWidth = 0.00276161

    def __init__(self):
        smach.State.__init__(
            self,
            outcomes=['fail','abort','success','request', 'center'],
            output_keys=['base_goal', 'picture_center', 'cue', 'cam_model'],
            input_keys=['faces', 'cam_model']
        )
        rospy.loginfo("Initialized Rule of Thirds state.")
        self.pose_pub = rospy.Publisher('/rules/rot_pose', PoseStamped)
        self.tf = tf.TransformListener()

    def calculateFov(self, cam):
        prosilica = rospy.get_param('use_prosilica', False)
        if prosilica:
            return self.calculateProcilicaFov(cam)
        else:
            return self.angleOfView

    def calculateProcilicaFov(self, cam):
        return 2 * atan(self.sensorWidth/(2*cam.fx()))

    def execute(self, data):
        if rospy.is_shutdown():
            return 'abort'
        if not data.faces:
            rospy.logerr(
                "Called to frame faces, but received NoneType for faces")
            return 'fail'
        if len(data.faces.faces) == 0:
            return 'fail'
        try:
            facePoints = [face.center for face in data.faces.faces]
            points = self.findPoints(facePoints, data.cam_model)
            point = self.findCloserPoint(data.faces.header.frame_id, points)
            framingCenter = self.findCenter(facePoints, data.faces.header.frame_id)
            angle = self.findAngle(point, framingCenter)
            goal  = MoveBaseGoal()
            goal.target_pose.header = data.faces.header
            goal.target_pose.pose.position = point
            goal.target_pose.pose.orientation = Quaternion(
                *tf.transformations.quaternion_from_euler(0,0,angle, 'rxyz'))
            data.base_goal = goal
            center = PointStamped()
            center.point = framingCenter
            center.header = data.faces.header
            data.picture_center = center
            self.pose_pub.publish(goal.target_pose)
            return 'success'
        except smach.InvalidUserCodeError as e:
            rospy.logerr("Failed to execute rule of thirds: %s" % e)
            return 'fail'
        except TooFarToFrame:
            data.cue = 'move_in'
            return 'request'
        except SubjectsTooClose:
            return 'center'

    def findCloserPoint(self, frame_id, points):
        dists = {}
        for point in points:
            ps = PointStamped()
            ps.point = point
            ps.header.frame_id = frame_id
            transformed = self.tf.transformPoint('/base_link', ps)
            dists[point] = transformed.point.x**2 + transformed.point.y**2
        return min(dists, key=lambda k: dists[k])
        

    def findAngle(self, robot, center):
        x = center.x - robot.x
        if x == 0:
            x = 0.00001
        y = center.y - robot.y
        if x<0:
            shift = pi
        else:
            shift = 0
        rospy.loginfo("(x,y, shift) = (%f, %f, %f)", x, y, shift)
        return shift + atan(y/x)

    def findPoints(self, points, cam):
        ends = self.findEnds(points)
        line = self.findOrthogonalLine(ends)
        dist = self.constrainToReasonablePhotoDistance(
            self.getDistanceFromSubject(line['point'], ends[0], cam)
        )
        result = self.findPointsOnLine(line,dist)
        return result

    def findCenter(self, points, frame_id):
        line = self.findOrthogonalLine(self.findEnds(points))
        ps = PointStamped()
        ps.point = line['point']
        ps.header.frame_id = frame_id
        transformed = self.tf.transformPoint('/head_tilt_link', ps)
        line['point'].z -= sqrt(transformed.point.y**2+transformed.point.x**2)/sqrt(3)
        return line['point']


    def findPointsOnLine(self, line, dist):
        m = line['slope']
        o = line['point']
        b = o.y-m*o.x
        x = (o.x + dist/sqrt(1+m**2),o.x - dist/sqrt(1+m**2))
        y = (m*x[0]+b,m*x[1]+b)

        points = (Point(x[0],y[0],0), Point(x[1], y[1], 0))
        return points

    def findOrthogonalLine(self, points):
        p1, p2 = points
        x,y,z  = (p1.x+p2.x)/2, (p1.y+p2.y)/2, (p1.z+p2.z)/2
        d = p2.x-p1.x
        if d == 0:
            rospy.loginfo("Guarding against zero when finding orthog. line.")
            d = 0.00001
        slope  = (p2.y-p1.y)/d
        if slope != 0:
            slope = -1/slope
        return {'point': Point(x,y,z), 'slope':slope}

    def constrainToReasonablePhotoDistance(self, dist):
        rospy.loginfo("Original distance is %s" % dist)
        if dist < self.minDistance:
            raise SubjectsTooClose
        if dist > self.maxDistance:
            raise TooFarToFrame()
        return dist

    def getDistanceFromSubject(self, center, subject, cam):
        # half of angle is for one side, get 1/3 from there
        return self.distance(center, subject)/tan(self.calculateFov(cam)/6.0)

    def closestPointOnLine(self, line, point):
        b2 = point.y+line['slope']*point.x
        b1 = line['point'].y-line['point'].x*line['slope']
        x  = (b2-b1)/(2*max(0.01,line['slope']))
        y  = -line['slope']*x+b2
        return Point(x,y,0)

    def distance(self, point, otherPoint):
        return sqrt((point.x-otherPoint.x)**2+(point.y-otherPoint.y)**2)

    def findEnds(self, points):
        maxDistance = -float('infinity')
        if len(points) < 2:
            return (points[0],points[0])
        for point in points:
            for otherPoint in points:
                if point != otherPoint:
                    dist = self.distance(point, otherPoint)
                    if maxDistance < dist:
                        bestPoints  = (point, otherPoint)
                        maxDistance = dist
        return bestPoints
